﻿using LDASharp;
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace GibbsLDASharp.LDA
{
    public class LDADataset
    {
        #region Instance Variables

        public LDADictionary localDict;			// local dictionary	
        public Document[] docs; 		// a list of documents	
        public int M; 			 		// number of documents
        public int V;			 		// number of words

        /// <summary>
        /// map from local coordinates (id) to global ones, null if the global dictionary is not set
        /// </summary>
        public Dictionary<int, int> lid2gid;

        /// <summary>
        /// link to a global dictionary (optional), null for train data, not null for test data
        /// </summary>
        public LDADictionary globalDict;
        #endregion

        #region Constructor
        /// <summary>
        /// 构造函数，初始化变量
        /// </summary>
        public LDADataset()
        {
            localDict = new LDADictionary();
            M = 0;
            V = 0;
            docs = null;

            globalDict = null;
            lid2gid = null;
        }

        public LDADataset(int M)
        {
            localDict = new LDADictionary();
            this.M = M;
            this.V = 0;
            docs = new Document[M];

            globalDict = null;
            lid2gid = null;
        }

        public LDADataset(int M, LDADictionary globalDict)
        {
            localDict = new LDADictionary();
            this.M = M;
            this.V = 0;
            docs = new Document[M];

            this.globalDict = globalDict;
            lid2gid = new Dictionary<int, int>();
        }
        #endregion

        /// <summary>
        /// set the document at the index idx if idx is greater than 0 and less than M
        /// </summary>
        /// <param name="doc">document to be set</param>
        /// <param name="idx">index in the document array</param>
        public void SetDoc(Document doc, int idx)
        {
            if (0 <= idx && idx < M)
            {
                docs[idx] = doc;
            }
        }
        /// <summary>
        /// set the document at the index idx if idx is greater than 0 and less than M
        /// </summary>
        /// <param name="str">string contains doc</param>
        /// <param name="idx">index in the document array</param>
        public void SetDoc(String str, int idx)
        {
            if (0 <= idx && idx < M)
            {
                String[] words = str.Split(new char[] { '[', ' ', '\t', '\n', ']' });

                List<int> ids = new List<int>();

                foreach (string word in words)
                {
                    int _id = localDict.word2id.Count;

                    if (localDict.Contains(word))
                        _id = localDict.GetID(word);

                    if (globalDict != null)
                    {
                        //get the global id		
                        try
                        {
                            int id = globalDict.GetID(word);
                            //Console.WriteLine(id);
                            localDict.AddWord(word);

                            lid2gid.Add(_id, id);
                            ids.Add(_id);
                        }
                        catch
                        { //not in global dictionary
                            //do nothing currently
                        }
                    }
                    else
                    {
                        localDict.AddWord(word);
                        ids.Add(_id);
                    }
                }

                Document doc = new Document(ids, str);
                docs[idx] = doc;
                V = localDict.word2id.Count;
            }
        }

        #region I/O methods


        /// <summary>
        /// read a dataset from a stream, create new dictionary
        /// </summary>
        /// <param name="filename">file name</param>
        /// <returns>dataset if success and null otherwise</returns>
        public static LDADataset ReadDataSet(String filename)
        {
            try
            {
                StreamReader reader = new StreamReader(filename, Encoding.UTF8);

                LDADataset data = ReadDataSet(reader);

                reader.Close();
                return data;
            }
            catch (Exception e)
            {
                Console.WriteLine("Read Dataset Error: " + e.Message);
                Console.WriteLine(e.StackTrace);
                return null;
            }
        }

        /// <summary>
        /// read a dataset from a file with a preknown vocabulary
        /// </summary>
        /// <param name="filename">read a dataset from a file with a preknown vocabulary</param>
        /// <param name="dict">the dictionary</param>
        /// <returns>dataset if success and null otherwise</returns>
        public static LDADataset ReadDataSet(String filename, LDADictionary dict)
        {
            try
            {
                StreamReader reader = new StreamReader(filename, Encoding.UTF8);
                LDADataset data = ReadDataSet(reader, dict);

                reader.Close();
                return data;
            }
            catch (Exception e)
            {
                Console.WriteLine("Read Dataset Error: " + e.Message);
                Console.WriteLine(e.StackTrace); ;
                return null;
            }
        }

        /// <summary>
        /// read a dataset from a stream, create new dictionary
        /// </summary>
        /// <param name="reader">stream from which we read dataset</param>
        /// <returns>dataset if success and null otherwise</returns>
        public static LDADataset ReadDataSet(StreamReader reader)
        {
            try
            {
                //read number of document
                String line;
                line = reader.ReadLine();
                int M = int.Parse(line);

                LDADataset data = new LDADataset(M);
                for (int i = 0; i < M; ++i)
                {
                    line = reader.ReadLine();

                    data.SetDoc(line, i);
                }

                return data;
            }
            catch (Exception e)
            {
                Console.WriteLine("Read Dataset Error: " + e.Message);
                Console.WriteLine(e.StackTrace);
                return null;
            }
        }

        /// <summary>
        /// Read a dataset from a stream with respect to a specified dictionary
        /// </summary>
        /// <param name="reader">stream from which we read dataset</param>
        /// <param name="dict">the dictionary</param>
        /// <returns>dataset if success and null otherwise</returns>
        public static LDADataset ReadDataSet(StreamReader reader, LDADictionary dict)
        {
            try
            {
                //read number of document
                String line;
                line = reader.ReadLine();
                int M = int.Parse(line);
                Console.WriteLine("NewM:" + M);

                LDADataset data = new LDADataset(M, dict);
                for (int i = 0; i < M; ++i)
                {
                    line = reader.ReadLine();

                    data.SetDoc(line, i);
                }

                return data;
            }
            catch (Exception e)
            {
                Console.WriteLine("Read Dataset Error: " + e.Message);
                Console.WriteLine(e.StackTrace);
                return null;
            }
        }

        /// <summary>
        /// read a dataset from a string, create new dictionary
        /// </summary>
        /// <param name="strs">String from which we get the dataset, documents are seperated by newline character </param>
        /// <returns>dataset if success and null otherwise</returns>
        public static LDADataset readDataSet(String[] strs)
        {
            LDADataset data = new LDADataset(strs.Count());

            for (int i = 0; i < strs.Count(); ++i)
            {
                data.SetDoc(strs[i], i);
            }
            return data;
        }


        /// <summary>
        /// read a dataset from a string with respect to a specified dictionary
        /// </summary>
        /// <param name="strs">String from which we get the dataset, documents are seperated by newline character</param>
        /// <param name="dict">the dictionary</param>
        /// <returns>dataset if success and null otherwise</returns>
        public static LDADataset ReadDataSet(String[] strs, LDADictionary dict)
        {
            //Console.WriteLine("readDataset...");
            LDADataset data = new LDADataset(strs.Count(), dict);

            for (int i = 0; i < strs.Count(); ++i)
            {
                //Console.WriteLine("set doc " + i);
                data.SetDoc(strs[i], i);
            }
            return data;
        }

        #endregion

        #region Non-file
        public static LDADataset ReadDataSet(string[] articles)
        {
            try
            {
                //read number of document
                string line;

                int M = articles.Length;

                LDADataset data = new LDADataset(M);
                for (int i = 0; i < M; ++i)
                {
                    line = articles[i];
                    data.SetDoc(line, i);
                }

                return data;
            }
            catch (Exception e)
            {
                Console.WriteLine("Read Dataset Error: " + e.Message);
                Console.WriteLine(e.StackTrace);
                return null;
            }
        }

        public static LDADataset ReadDataSet(IEnumerable<string> articles)
        {
            try
            {
                int M = articles.Count();

                LDADataset data = new LDADataset(M);
                int i = 0;
                foreach (string line in articles)
                {

                    data.SetDoc(line, i);
                    i++;
                }

                return data;
            }
            catch (Exception e)
            {
                Console.WriteLine("Read Dataset Error: " + e.Message);
                Console.WriteLine(e.StackTrace);
                return null;
            }
        }
        #endregion
    }
}
